#include "../include/SSS_inter.h"
#include <thrust/reduce.h>
#include <thrust/execution_policy.h>
#include <thrust/device_ptr.h>
#include <thrust/device_vector.h>




// #define gridsize 65535
// #define blocksize 256

#define zero_tol 1e-16

#define block_size 64


double get_time(void)
{
    struct timeval tv;
    double t;

    gettimeofday(&tv,NULL);
    t = tv.tv_sec * 1000 + tv.tv_usec /1000;
    return t;
}




void __global__ jacobi_prolongation_check_dig_kernel(int row,double *d_dig)
{
    int tid =blockIdx.x*blockDim.x+ threadIdx.x;

    if(tid < row)    
    {
        if (SSS_ABS(d_dig[tid]) < 1e-4) 
        d_dig[tid] = 1.0;
    }

}


void __global__ jacobi_prolongation_jacobi_kernel(int row,double smooth_factor, double *d_dig, int *S_row_ptr , int *S_col_idx, double *S_val,double *A_val)
{
    int tid =blockIdx.x*blockDim.x+ threadIdx.x;

//    printf("row = %d ,smooth_factor = %lf\n",row,smooth_factor);

    if(tid < row)    
    {
        //1 2 
        for (int j=S_row_ptr[tid]; j<S_row_ptr[tid+1]; ++j) 
        {
            if (S_col_idx[j] == tid) 
            {
                S_val[j] = 1 - smooth_factor * A_val[j] / d_dig[tid];
            }
            else 
            {
                S_val[j] = - smooth_factor * A_val[j] / d_dig[tid];
            }
        }

    }

}

// void __global__ T_get_nnz_kernel(int row,int *aggregates,int *row_nnz)
// {
//     int tid = blockDim.x * blockIdx.x +threadIdx.x;
    
//     if(tid < row)
//     {
//         int num = 0;
//         if(aggregates[tid] > UNPT) 
//         {
//             num++;
//         }
//         row_nnz[tid] = num;
//     }
// }



void __global__ T_struct_kernel(int row,int *aggregates,int *row_ptr,int *col_idx,double *val,double *B)
{  
    int tid = blockDim.x * blockIdx.x +threadIdx.x;

    if(tid < row)
    {
            row_ptr[tid] = tid;
            col_idx[tid] = aggregates[tid];
            val[tid] = B[tid];
    }
    if (tid == row)
    {
        row_ptr[tid] = tid;
    }

}

SSS_MAT tentative_p_GPU(SSS_IVEC *aggregates,SSS_VEC *B,SSS_AMG_PARS * pars)
{
    cudaError EE;
 
    printf("form_tentative_p form_tentative_p\n");
    SSS_MAT T;
    
    /* Form tentative prolongation */
    T.num_rows = aggregates->n;
    T.num_cols = pars->n_agg;

    const int row = T.num_rows;

    int grid_size = ceil ((double)row / (double)block_size);

    T.num_nnzs = row;

    //GPU
    T = SSS_mat_struct_create_device(T.num_rows,T.num_cols,T.num_nnzs,T);


    printf("T.row = %d \n",row);

    T_struct_kernel<<<grid_size,block_size>>>(row,aggregates->d_d,T.d_row_ptr,T.d_col_idx,T.d_val,B->d_d);

    cudaDeviceSynchronize();    


    //T
    int *T_idx1 = (int*)malloc(sizeof(int) * (T.num_rows+1));
    cudaMemcpy(T_idx1,T.d_row_ptr, (T.num_rows+1) *sizeof(int), cudaMemcpyDeviceToHost);

    int *T_idx2 = (int*)malloc(sizeof(int) *T.num_nnzs);
    cudaMemcpy(T_idx2,T.d_col_idx, (T.num_nnzs) *sizeof(int), cudaMemcpyDeviceToHost);

    double *T_val = (double *)malloc(sizeof(double ) * T.num_nnzs);
    cudaMemcpy(T_val,T.d_val, (T.num_nnzs) *sizeof(double), cudaMemcpyDeviceToHost);

    for(int i=0;i<10;i++)
    {
        printf("T:row_ptr[%d] = %d,col_idx[%d] = %d, val[%d] = %lf\n",i,T_idx1[i],i,T_idx2[i],i,T_val[i]);    
    }


    return T;
}







SSS_MAT SSS_amg_jacobi_prolongation_GPU(SSS_MAT* A, SSS_MAT* T, SSS_MAT *P)
{
    cudaError_t EE;
    double prolong_weight = 4.0 / 3;
    int   row = A->num_rows, col= A->num_cols;
    double  smooth_factor = prolong_weight;
    
    smooth_factor = 0.67;

    //GPU S
    SSS_MAT S = SSS_mat_struct_create_device(A->num_rows, A->num_cols,A->num_nnzs,S);

    double row_sum_A, row_sum_N;
    int i,j;

    printf("step 1 \n");

    EE = cudaMemcpy(S.d_row_ptr,A->d_row_ptr, (row+1)* sizeof(int), cudaMemcpyDeviceToDevice);
    if(EE != cudaSuccess)
    {
        printf("d_row_ptr fail \n");
    }
    EE = cudaMemcpy(S.d_col_idx,A->d_col_idx, A->num_nnzs* sizeof(int), cudaMemcpyDeviceToDevice);
    if(EE != cudaSuccess)
    {
        printf("d_col_idx fail \n");
    }
    EE = cudaMemcpy(S.d_val,A->d_val, A->num_nnzs* sizeof(double), cudaMemcpyDeviceToDevice);
    if(EE != cudaSuccess)
    {
        printf("d_col_idx fail \n");
    }

    printf("step 1.6 \n");

    int grid_size = ceil ((double)row / (double)block_size);
    jacobi_prolongation_check_dig_kernel<<<grid_size,block_size>>>(row,A->d_dig);
        
    cudaDeviceSynchronize();    

    jacobi_prolongation_jacobi_kernel<<<grid_size,block_size>>>(row,smooth_factor,A->d_dig,S.d_row_ptr,S.d_col_idx,S.d_val,A->d_val);
    
    cudaDeviceSynchronize();    


    printf("step 2 \n");

    //P的device空间以创建好

    printf("1111111111111111111\n");
    
    //S
    int *S_idx1 = (int*)malloc(sizeof(int) * (S.num_rows+1));
    cudaMemcpy(S_idx1,S.d_row_ptr, (S.num_rows+1) *sizeof(int), cudaMemcpyDeviceToHost);

    int *S_idx2 = (int*)malloc(sizeof(int) *S.num_nnzs);
    cudaMemcpy(S_idx2,S.d_col_idx, (S.num_nnzs) *sizeof(int), cudaMemcpyDeviceToHost);

    double *S_val = (double *)malloc(sizeof(double ) * S.num_nnzs);
    cudaMemcpy(S_val,S.d_val, (S.num_nnzs) *sizeof(double), cudaMemcpyDeviceToHost);

    for(int i=0;i<10;i++)
    {
        printf("S:row_ptr[%d] = %d,col_idx[%d] = %d, val[%d] = %lf\n",i,S_idx1[i],i,S_idx2[i],i,S_val[i]);    
    }


    *P = SSS_blas_spgemm_device(&S,T,P);
    
    
    
    printf("22222222222222222222222\n");
    return *P;

}